package org.gilgamesh.stool.search.es.util;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.StrUtil;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.index.query.BoolQueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.index.query.RangeQueryBuilder;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.SearchHits;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.gilgamesh.stool.search.constant.RangeType;
import org.gilgamesh.stool.search.es.ESSearchRequest;
import org.gilgamesh.stool.search.es.ESSearchResponse;
import org.gilgamesh.stool.search.es.param.BoolParam;
import org.gilgamesh.stool.search.es.param.RangeTimeParam;
import org.gilgamesh.stool.search.es.result.Result;

import java.io.IOException;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 查询工具类
 *
 * @author Gilgamesh
 * @since 2021-08-20
 **/
public class SearchUtil {

    public static ESSearchResponse search(ESSearchRequest request, RestHighLevelClient esClient) {
        SearchRequest searchRequest = new SearchRequest();
        /* 设置索引名称 */
        searchRequest.indices(request.indexName());
        SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
        /* 设置需要返回的字段，和不需要返回的字段 */
        sourceBuilder.fetchSource(request.needFields(), request.unNeedFields());

        BoolQueryBuilder boolQuery = QueryBuilders.boolQuery();
        /* 设置must条件 */
        if (request.hasMust()) {
            List<BoolParam> musts = request.must();
            for (BoolParam must : musts) {
                List<Object> values = must.getValues();
                for (Object value : values) {
                    boolQuery.must(QueryBuilders.matchQuery(must.getField(), value));
                }
            }
        }
        /* 设置should条件 */
        if (request.hasShould()) {
            List<BoolParam> shoulds = request.should();
            for (BoolParam should : shoulds) {
                List<Object> values = should.getValues();
                for (Object value : values) {
                    boolQuery.should(QueryBuilders.matchQuery(should.getField(), value));
                }
            }
            boolQuery.minimumShouldMatch(1);
        }

        /* 请求中设置了按照时间排序 */
        if (request.isRangeTime()) {
            RangeTimeParam rangeTime = request.rangeTime();
            String rangeField = rangeTime.field();
            /* 大于 */
            RangeQueryBuilder rangeQueryBuilder = QueryBuilders.rangeQuery(rangeField);
            if (rangeTime.rangeMap().containsKey(RangeType.GT)) {
                Long gt = rangeTime.rangeMap().get(RangeType.GT);
                rangeQueryBuilder.gt(gt);
            }
            /* 小于 */
            if (rangeTime.rangeMap().containsKey(RangeType.LT)) {
                Long lt = rangeTime.rangeMap().get(RangeType.LT);
                rangeQueryBuilder.lt(lt);
            }
            /* 大于等于 */
            if (rangeTime.rangeMap().containsKey(RangeType.GTE)) {
                Long gte = rangeTime.rangeMap().get(RangeType.GTE);
                rangeQueryBuilder.gte(gte);
            }
            /* 小于等于 */
            if (rangeTime.rangeMap().containsKey(RangeType.LTE)) {
                Long lte = rangeTime.rangeMap().get(RangeType.LTE);
                rangeQueryBuilder.lte(lte);
            }
            boolQuery.filter(rangeQueryBuilder);
        }

        /* 排序 */
        if (request.isSorted()) {
            sourceBuilder.sort(request.sortField(), request.sort());
        }
        /* 想要的条数 */
        if (request.size() > 0) {
            sourceBuilder.size(request.size());
        }
        sourceBuilder.query(boolQuery);
        searchRequest.source(sourceBuilder);

        ESSearchResponse eSResponse = new ESSearchResponse();
        try {
            SearchResponse response = esClient.search(searchRequest, RequestOptions.DEFAULT);
            SearchHits hits = response.getHits();
            /* 设置总查询条件 */
            eSResponse.total((int) hits.getTotalHits().value);
            Map<String, Result> resultMap = new HashMap<>();
            for (SearchHit hit : hits) {
                Map<String, Object> sourceAsMap = hit.getSourceAsMap();
                String key = request.choosePrimaryKeyField();
                for (String resKey : sourceAsMap.keySet()) {
                    if(sourceAsMap.containsKey(key)){
                        Result result = new Result();
                        result.setK(resKey);
                        result.setV(sourceAsMap.get(resKey).toString());
                        resultMap.put(key,result);
                    }
                }
            }
            return eSResponse.resultMap(resultMap);
        } catch (IOException e) {
            e.printStackTrace();
        }
        return null;
    }
}
